package course.p7_binaryTree.s5_RedBlockTree;

import course.p7_binaryTree.Node;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

public class RBTree {

    Node root;


    /**
     * 在节点的左孩子的左子树插入
     * 右旋操作,操作步骤是：
     * 首先得到root的左孩子left
     * 然后将root的左孩子设置left的右孩子
     * 然后将left的右孩子设置为root
     * @param root
     */
    void rightRotate(Node root) {
        // 得到root的左节点
        Node left = root.getLeft();
        // 将root的左子树设置为left的右子树
        root.setLeft(left.getRight());
        // 将left的右节点设为root
        left.setRight(root);
        // 父引用维护
        left.parent = root.parent;
        if(root.parent!=null) {
            if(root == root.parent.left){
                root.parent.left = left;
            }
            else{
                root.parent.right = left;
            }
        }else this.root = left;
        root.parent = left;
    }


    /**
     * 在右孩子的右子树插入元素
     * 左旋操作,操作步骤是：
     * 首先得到root的右孩子right
     * 然后将root的右孩子设置为right的左孩子
     * 最后将right的左孩子设置为root
     * @param root
     */
    void leftRotate(Node root) {
        Node right = root.getRight();
        root.setRight(right.getLeft());
        right.setLeft(root);
        // 父引用维护
        right.parent = root.parent;
        if(root.parent!=null) {
            if(root == root.parent.left){
                root.parent.left = right;
            }else{
                root.parent.right = right;
            }
        }else this.root = right;
        root.parent = right;
    }


    /**
     * 红黑树的插入操作的步骤如下：
     * 1. 根据比较找出预插入节点的父节点
     * 2. 进行特殊情况处理(首次插入)
     * 3. 正常插入
     * 4. 维护插入节点的红黑特性
     * @param data 插入的数据
     * @return
     */
    Node insert(int data) {
        //z 需要插入的节点
        Node z = new Node(data);
        Node cur = root;
        Node parent = null;
        while (cur!=null) {
            parent = cur;
            if(z.data < cur.data) {
                cur = cur.left;
            }
            else if(z.data>cur.data){
                cur = cur.right;
            }
            else return null;
        }
        z.parent = parent;
        //特殊情况处理，首次插入
        if(parent == null) {
            this.root = z;
            z.color = false;
            return this.root;
        }
        else if(z.data < parent.data){
            parent.left = z;
        }
        else {
            parent.right = z;
        }
        z.color = true;
        insertFix(z);
        return z;
    }

    /**
     * 红黑树插入修复操作
     * @param z
     */
    void insertFix(Node z) {
        while (true){
             if(z.parent!=null&& z.parent.color) {
                    Node uncle = null;
                    if(z.parent.parent.left == z.parent){
                        uncle = z.parent.parent.right;
                        if(uncle==null||!uncle.color){
                            //情况2
                            if(z==z.parent.right) {
                                z = z.parent;
                                leftRotate(z);
                                continue;
                            }
                            //情况3
                            if(z == z.parent.left) {
                                z.parent.color = false;
                                z.parent.parent.color = true;
                                rightRotate(z.parent.parent);
                                continue;
                            }
                        }
                        //情况1
                        else {
                            z.parent.color = false;
                            uncle.color = false;
                            if(z.parent.parent!=root){
                                z.parent.parent.color = !z.parent.parent.color;
                            }
                            z = z.parent.parent;
                            continue;
                        }
                    }
                    else {
                        uncle = z.parent.parent.left;
                        if(uncle==null||!uncle.color) {
                            //情况5
                            if(z==z.parent.left) {
                                z = z.parent;
                                rightRotate(z);
                                continue;
                            }
                            //情况6
                            if(z == z.parent.right) {
                                z.parent.color = false;
                                z.parent.parent.color = true;
                                leftRotate(z.parent.parent);
                                continue;
                            }
                        }
                        //情况4
                        else {
                            z.parent.color = false;
                            uncle.color = false;
                            if(z.parent.parent!=root){
                                z.parent.parent.color = !z.parent.parent.color;
                            }
                            z = z.parent.parent;
                            continue;
                        }
                    }
            }
            break;
        }
        this.root.color = false;
    }

    /**
     * 将节点v替换u节点的位置
     * @param u
     * @param v
     */
    public void RBTransplant(Node u,Node v){
        if(u.parent == null){
            root = v;
        }
        else if(u == u.parent.left){
            u.parent.left = v;
        }
        else{
            u.parent.right = v;
        }
        if(v!=null)
           v.parent = u.parent;
    }

    /**
     * 删除节点z
     * @param z 要删除的节点
     */
    public void RBDelete(Node z) {
        Node y = z, x = null;
        boolean yColor = y.color;
        // 如果z的左子树为空
        // 则直接将z的右孩子替换掉z即可
        if(z.left == null){
            x = z.right;
            RBTransplant(z,z.right);
        }
        // 同理如果z的右子树为空，将其左孩子替换掉它
        else if(z.right == null) {
            x = z.left;
            RBTransplant(z,z.left);
        }
        // 否则进行普通是删除操作
        else {
            // 找到z右子树的最小节点
            y = findMin(z.right);
            yColor = y.color;
            x = y.right;
            if(y.parent == z) {
                x.parent = y;
            }
            else {
                RBTransplant(y,y.right);
                y.right = z.right;
                y.right.parent = y;
            }
            RBTransplant(z,y);
            y.left = z.left;
            y.left.parent = y;
            y.color = z.color;
        }
        if(!yColor&&x!=null){
            RBDeleteFix(x);
        }
    }

    /**
     * 删除后的修复工作
     * @param x
     */
    public void RBDeleteFix(Node x) {
        while (x!=root&&!x.color){
            Node w;
            if(x == x.parent.left) {
                w = x.parent.right;
                if(w.color){
                    w.color = false;
                    x.parent.color = true;
                    leftRotate(x.parent);
                    w = x.parent.right;
                }
                if(!w.left.color&&!w.right.color){
                    x = x.parent;
                }
                else if(!w.right.color) {
                    w.left.color = false;
                    w.color = true;
                    rightRotate(w);
                    w = x.parent.right;
                }
                w.color = x.parent.color;
                x.parent.color = false;
                w.right.color = false;
                leftRotate(x.parent);
            }
            else {
                //将前面的操作中的left和right互换
                w = x.parent.left;
                if(w.color){
                    w.color = false;
                    x.parent.color = true;
                    rightRotate(x.parent);
                    w = x.parent.left;
                }
                if(!w.right.color&&!w.left.color){
                    x = x.parent;
                }
                else if(!w.left.color) {
                    w.right.color = false;
                    w.color = true;
                    leftRotate(w);
                    w = x.parent.left;
                }
                w.color = x.parent.color;
                x.parent.color = false;
                w.left.color = false;
                rightRotate(x.parent);
            }
            x = root;
            x.color = false;
        }
    }


    /**
     * 寻找树中最小元素
     * 基本思想是此树的最小元素一定在最左边
     * @param root
     * @return
     */
    Node findMin(Node root) {
        if(root==null) return null;
        while(root.getLeft()!=null)
            root = root.getLeft();
        return root;
    }

    /**
     * 中序遍历
     * @param first
     */
    void Check3(Node first) {
        if(first!=null) {
            Check3(first.getLeft());
            System.out.print(first.getData());
            if(first.color){
                System.out.print("(red) ");
            }else{
                System.out.print("(black) ");
            }
            Check3((first.getRight()));
        }
    }

    /**
     * 将层次遍历的结果存入二维数组中
     * @param root
     * @return
     */
    public List<List<Integer>> levelOrder(Node root) {
        Queue<Node> queue= new LinkedList<>();
        List<List<Integer>> result = new ArrayList<>();
        if(root!=null) queue.add(root);
        while(!queue.isEmpty()) {
            List<Integer> temp = new ArrayList<>();
            int size = queue.size();
            for(int i=0;i<size;i++) {
                Node node = queue.poll();
                assert node != null;
                temp.add(node.getData());
                if(node.getLeft()!=null) queue.add(node.getLeft());
                if(node.getRight()!=null) queue.add(node.getRight());
            }
            result.add(temp);
        }
        return result;
    }


}
